"""

1.使用torchvision来加载和标准化CIFAR10训练和测试数据集
2.使用pytorch框架定义一个卷积神经网络CNN
3.定义一个损失函数
4.在训练集上训练网络
5.在测试数据集上测试网络
6.在不同的类上测试网络

"""

import torch
import torchvision
import torchvision.transforms as transforms
import torch.utils.data.dataloader
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np

import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

#totensor:导入的数据是PILImage格式的，需要转换成tensor的格式
"""
前面的（0.5，0.5，0.5） 是 RGB三个通道上的均值，后面(0.5, 0.5, 0.5)是三个通道的标准差，
Normalize对每个通道执行以下操作：image =（图像-平均值）/ std在您的情况下，参数mean，std分别以0.5和0.5的形式传递。
这将使图像在[-1,1]范围内归一化。
"""
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])

train_data = torchvision.datasets.CIFAR10(root='./CIFAR10data',train=True,download=False,transform=transform)
train_loader = torch.utils.data.DataLoader(train_data,batch_size=5,shuffle=True,num_workers=0)

test_data = torchvision.datasets.CIFAR10(root='./CIFAR10data',train=False,download=False,transform=transform)
test_loader = torch.utils.data.DataLoader(test_data,batch_size=5,shuffle=False,num_workers=0)

classes = ('plane','car','bird','cat','deer','dog','frog','horse','ship','trunk')
def imgshow(img):
    img = img/2+0.5 #unnormalize
    npimg = img.numpy()
    # np.transpose:按需求转置
    plt.imshow(np.transpose(npimg,(1,2,0)))
#CIFAR-10数据集由10个类的60000个32x32彩色图像组成，每个类有6000个图像。有50000个训练图像和10000个测试图像。
class Net1(nn.Module):
    def __init__(self):
        super(Net1, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 12, 5)
        self.fc1 = nn.Linear(12 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 60)
        self.fc3 = nn.Linear(60, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 12 * 5 * 5)  # 拉成向量
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        return x


net = Net1()
criterion = nn.CrossEntropyLoss()
optimzer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
maxx =999
# 训练网络
for epoch in range(20):
    running_loss = 0.0
    for i, data in enumerate(train_loader, 0):
        input, target = data
        input, target = Variable(input), Variable(target)
        input, target = input.to(device), target.to(device)

        output = net(input)
        loss = criterion(output, target)
        optimzer.zero_grad()
        loss.backward()
        optimzer.step()
        running_loss += loss.item()
        if i % 2000 == 1999:
            print('[%d,%5d]loss:%.3f' % (epoch + 1, i + 1, running_loss / 2000))
            if running_loss / 2000 < maxx:
                maxx = running_loss / 2000
                torch.save(net, 'MNIST.pkl')
            running_loss = 0.0
print('Finished Training')
print(maxx)
# 保存训练模型
#torch.save(net, 'MNIST.pkl')
net = torch.load('MNIST.pkl')

dataiter = iter(test_loader)
images, labels = dataiter.next()
imgshow(torchvision.utils.make_grid(images))
print('GroundTruth:', ''.join('%5s' % classes[labels[j]] for j in range(4)))
"""outputs = net(Variable(images))
_, pred = torch.max(outputs.data, 1)
print('Predicted:', ''.join('%5s' % classes[pred[j].item()] for j in range(4)))"""
correct = 0.0
total = 0
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = torch.load('./MNIST.pkl')
for data in test_loader:
    images,labels = data
    images, labels = images.to(device), labels.to(device)
    outputs = net(Variable(images))
    _, pred = torch.max(outputs.data, 1)
    total += labels.size(0)
    correct += (pred==labels).sum()
print('Accuracy of the network on the 10000 test images:%d %%'%(100*correct/total))

